import numpy as np
from scipy import signal
from scipy.fft import fft, fftfreq
 
def calculate_power_high_precision(input_signal, fs, signal_freq, n_fft_mult=4):
    N = len(input_signal)
    input_signal = input_signal - np.mean(input_signal)  # 去除直流
    
    # 窗函数生成与补偿
    window = signal.windows.flattop(N)
    coherent_gain = np.mean(window)
    window_power = np.sum(window ** 2) / N  # 归一化窗能量
    
    # 加窗并FFT（补零）
    n_fft = N * n_fft_mult
    spectrum = fft(input_signal * window, n=n_fft) / (N * coherent_gain)
    spectrum_abs = np.abs(spectrum[:n_fft//2]) * 2
    freq_axis = fftfreq(n_fft, 1/fs)[:n_fft//2]
    
    # 三阶多项式插值（扩展范围）
    peak_idx = np.argmin(np.abs(freq_axis - signal_freq))
    idx_low = max(0, peak_idx-3)
    idx_high = min(len(freq_axis), peak_idx+4)
    poly = np.polyfit(freq_axis[idx_low:idx_high], spectrum_abs[idx_low:idx_high], 3)
    true_peak = np.polyval(poly, signal_freq)
    
    return (true_peak ** 2) / 2  # 正弦波功率公式

# 测试
fs, f0, N = 1000, 50.3, 4096
t = np.arange(N) / fs
test_signal = 1.0 * np.sin(2 * np.pi * f0 * t)

power = calculate_power_high_precision(test_signal, fs, f0, n_fft_mult=4)
print(f"优化后功率: {power:.6f} (理论值: 0.5, 误差: {(power-0.5)/0.5*100:.4f}%)")